# Copyright (c) Phigent Robotics. All rights reserved.
from .bevdet import BEVDet
import torch.nn.functional as F
import torch
from mmdet.models import DETECTORS
from mmdet.models.builder import build_loss
from mmcv.cnn.bricks.conv_module import ConvModule
from torch import nn
import numpy as np
from ..import build_head
from ..import build_fusion_layer
from mmdet3d.models.utils import FFN
from mmdet3d.models.utils.spconv_voxelize import SPConvVoxelization
from mmcv.runner import force_fp32
import torch.nn.functional as F

@DETECTORS.register_module()
class LightFusionOCC(BEVDet):
    def __init__(self,
                 occ_head=None,
                 upsample=False,
                 fuse_neck=None,
                 **kwargs):
        super(LightFusionOCC, self).__init__(**kwargs)
        # 构建 occ_head 和 fuse_neck 模块
        self.occ_head_cfg = occ_head  
        self.occ_head = build_head(occ_head)
        self.fuse_neck = build_fusion_layer(fuse_neck)
        self.upsample = upsample

    @property
    def with_pts_bbox(self):
        return True

    def extract_feat(self, points, img, img_metas):
        """提取图像和点云特征"""
        img_feats = self.extract_img_feat(img, img_metas)
        pts_feats = self.extract_pts_feat(points, img_feats, img_metas)
        # 假设 pts_feats 为列表，这里取第一个
        return (img_feats, pts_feats[0])
    
    def extract_img_feat(self, img, img_metas):
        """提取图像特征"""
        img = self.prepare_inputs(img)
        x, _ = self.image_encoder(img[0])
        return [x]
    
    def forward_train(self,
                      points=None,
                      img_metas=None,
                      gt_bboxes_3d=None,
                      gt_labels_3d=None,
                      gt_labels=None,
                      gt_bboxes=None,
                      img_inputs=None,
                      proposals=None,
                      gt_bboxes_ignore=None,
                      **kwargs):
        """
        训练时的前向传播：
         1. 提取图像与点云特征；
         2. 图像视角转换为 BEV 特征；
         3. 融合图像 BEV 与点云特征，计算融合损失（如果 fuse_neck 定义了损失函数）；
         4. 后续 BEV 特征编码与 occ_head 损失计算。
        """
        # 1. 提取图像和点云特征

        img_feats, pts_feats = self.extract_feat(points, img=img_inputs, img_metas=img_metas)

        # 2. 将图像特征转换为 BEV 特征
        img_feats_bev, _depth = self.img_view_transformer(
            img_feats + img_inputs[1:7],
            depth_from_lidar=kwargs['gt_depth']
        )
        
        # 3. 获取辅助信息
        voxel_semantics = kwargs['voxel_semantics']  # (B, Dx, Dy, Dz)
        mask_camera = kwargs['mask_camera']          # (B, Dx, Dy, Dz)
        
        # 4. 将点云特征调整至与图像 BEV 特征相同尺寸
        pts_feats = F.interpolate(pts_feats,
                                  size=(img_feats_bev.shape[2], img_feats_bev.shape[3]),
                                  mode='bilinear',
                                  align_corners=True)
        
        # 5. 融合多模态特征（fuse_neck 模块返回 (deformed_feature, mm_bev)）
        bev_features_before_fusion = [img_feats_bev, pts_feats]
        fused_features = self.fuse_neck(bev_features_before_fusion)
        losses = dict()
        if isinstance(fused_features, (list, tuple)) and len(fused_features) == 2:
            deformed_feature, mm_bev = fused_features
            # 计算融合损失（返回一个 torch.Tensor）
            loss_fuse = self.fuse_neck.calculate_loss(deformed_feature, mm_bev)
            losses.update({'loss_fuse': loss_fuse})
            # 选择其中一种作为后续 BEV 特征，通常使用 deformed_feature
            bev_features_fused = deformed_feature
        else:
            bev_features_fused = fused_features
        
        # 6. 对融合后的 BEV 特征进行编码
        bev_feature = self.img_bev_encoder_backbone(bev_features_fused)
        bev_feature = self.img_bev_encoder_neck(bev_feature)
        if self.upsample:
            bev_features_fused = F.interpolate(bev_features_fused,
                                               scale_factor=2,
                                               mode='bilinear',
                                               align_corners=True)
        # 7. occ_head 损失计算（occ_head.loss 返回字典）
        loss_occ_bev_feature = self.forward_occ_train(bev_feature, voxel_semantics, mask_camera)
        losses.update(loss_occ_bev_feature)
        return losses

    def forward_occ_train(self, bev_features_fused, voxel_semantics, mask_camera):
        """计算 occ_head 分支的损失"""
        outs = self.occ_head(bev_features_fused)
        loss_occ = self.occ_head.loss(
            outs,           # occ_pred: (B, Dx, Dy, Dz, n_cls)
            voxel_semantics,  # 标签 (B, Dx, Dy, Dz)
            mask_camera       # 掩码 (B, Dx, Dy, Dz)
        )
        return loss_occ

    def simple_test(self,
                    points,
                    img_metas,
                    img_inputs=None,
                    rescale=False,
                    **kwargs):
        """测试时前向传播，返回 occ_head 的预测结果"""
        img_feats, pts_feats = self.extract_feat(points, img=img_inputs, img_metas=img_metas)
        img_feats_bev, _depth = self.img_view_transformer(
            img_feats + img_inputs[1:7],
            depth_from_lidar=kwargs['gt_depth']
        )
        pts_feats = F.interpolate(pts_feats,
                                  size=(img_feats_bev.shape[2], img_feats_bev.shape[3]),
                                  mode='bilinear',
                                  align_corners=True)
        bev_features_before_fusion = [img_feats_bev, pts_feats]
        fused_features = self.fuse_neck(bev_features_before_fusion)
        if isinstance(fused_features, (list, tuple)) and len(fused_features) == 2:
            deformed_feature, _ = fused_features
            bev_features_fused = deformed_feature
        else:
            bev_features_fused = fused_features
        if self.upsample:
            bev_features_fused = F.interpolate(bev_features_fused,
                                               scale_factor=2,
                                               mode='bilinear',
                                               align_corners=True)
        bev_feature = self.img_bev_encoder_backbone(bev_features_fused)
        bev_feature = self.img_bev_encoder_neck(bev_feature)
        occ_list = self.simple_test_occ(bev_feature, img_metas=img_metas)
        return occ_list

    def simple_test_occ(self, bev_feature, img_metas=None):
        """获取 occ_head 的预测结果"""
        outs = self.occ_head(bev_feature)
        if not hasattr(self.occ_head, "get_occ_gpu"):
            occ_preds = self.occ_head.get_occ(outs, img_metas)
        else:
            occ_preds = self.occ_head.get_occ_gpu(outs, img_metas)
        return occ_preds

    def forward_dummy(self,
                      points=None,
                      img_metas=None,
                      img_inputs=None,
                      **kwargs):
        """用于计算模型参数量的 dummy forward，不计算实际损失"""
        img_feats, pts_feats = self.extract_feat(points, img=img_inputs, img_metas=img_metas)
        img_feats_bev, _depth = self.img_view_transformer(
            img_feats + img_inputs[1:7],
            depth_from_lidar=kwargs['gt_depth']
        )
        pts_feats = F.interpolate(pts_feats,
                                  size=(img_feats_bev.shape[2], img_feats_bev.shape[3]),
                                  mode='bilinear',
                                  align_corners=True)
        bev_features_before_fusion = [img_feats_bev, pts_feats]
        fused_features = self.fuse_neck(bev_features_before_fusion)
        if isinstance(fused_features, (list, tuple)) and len(fused_features) == 2:
            deformed_feature, _ = fused_features
            bev_features_fused = deformed_feature
        else:
            bev_features_fused = fused_features
        if self.upsample:
            bev_features_fused = F.interpolate(bev_features_fused,
                                               scale_factor=2,
                                               mode='bilinear',
                                               align_corners=True)
        bev_feature = self.img_bev_encoder_backbone(bev_features_fused)
        bev_feature = self.img_bev_encoder_neck(bev_feature)
        outs = self.occ_head(bev_feature)
        return outs